'''
$Revision: 62 $
$LastChangedDate: 2008-03-13 18:10:27 +0000 (Thu, 13 Mar 2008) $
'''

import unittest
import os, sys

from MomentClosure.GetMomentEquation import getMomentEquation
from MomentClosure.GetReactionInformation import GetStoichiometryDict
from MomentClosure.VariableDict import Variables
from swiginac import numeric
from MomentClosure.GlobalVariables import globalVariables


class TestMomentEquations(unittest.TestCase):
    '''Tests the moment closure equations with know cases.
    
        If they work, everything else should as well.
    '''
    def setUp(self):
        pass

    def testImmigrationDeath(self):
        Variables().initialise()
        self.__testImmigrationDeath('xml_files/imdeath.xml')
        #With local parameter
        self.__testImmigrationDeath('xml_files/smfsb/ch05-imdeath.xml')
        
        #Reactions of the same type
        self.__testImmigrationDeath('xml_files/imdeath2.xml')
    
    def __testImmigrationDeath(self, file_name):
        Variables().initialise()
        gdd = GetStoichiometryDict(sbml_file = file_name)
        reaction_list = gdd.getReactionList()
        
        l = Variables().getParameter('lambda')
        nu = Variables().getParameter('mu')
        
        mu1 = Variables().getMoment((1,))
        mu2 = Variables().getMoment((2,))
        mu3 = Variables().getMoment((3,)) 
        
        mom_eq = getMomentEquation((0,), reaction_list) 
        self.assertEqual(mom_eq, 0)  
               
        #Mean
        mom_eq = getMomentEquation((1,), reaction_list) 
        exact = l - nu*mu1
        diff = mom_eq - exact
        self.assertEqual(diff.expand(), 0)  

        #Mu[2]
        mom_eq = getMomentEquation((2,), reaction_list) 
        exact = l*(1+2*mu1)+nu*(mu1-2*mu2)
        diff = mom_eq - exact
        self.assertEqual(diff.expand(), 0)  
        
        #Mu[3]        
        mom_eq = getMomentEquation((3,), reaction_list)
        exact = l*(1+3*mu1+3*mu2) - nu*(mu1-3*mu2+3*mu3)
        diff = mom_eq - exact
        self.assertEqual(diff.expand(), 0) 


    def testDimerModel(self):
        Variables().initialise()
        
        gdd = GetStoichiometryDict(sbml_file = 'xml_files/smfsb/single-dimer.xml')
        reaction_list = gdd.getReactionList()
        
        k1 =Variables().getParameter('k1')
        k2 = Variables().getParameter('k2')
        P0 = Variables().getParameter('P0')
        
        mu1 = Variables().getMoment((1,))
        mu2 = Variables().getMoment((2,))
        mu3 = Variables().getMoment((3,)) 
        mu4 = Variables().getMoment((4,))
        
        mom_eq = getMomentEquation((0,), reaction_list) 
        self.assertEqual(mom_eq, 0)
                  
        mom_eq = getMomentEquation((1,), reaction_list) 
        exact = k1*(mu1-mu2)+k2*(P0-mu1)
        diff = mom_eq - exact
        self.assertEqual(diff.expand(), 0)  
        
        mom_eq = getMomentEquation((2,), reaction_list) 
        exact = -2*k1*(mu1-2*mu2+mu3)+2*k2*(P0*(mu1+1)-mu1-mu2)
        diff = mom_eq - exact
        self.assertEqual(diff.expand(), 0)  

        mom_eq = getMomentEquation((3,), reaction_list) 
        exact = k1*(4*mu1-10*mu2+9*mu3-3*mu4)
        exact += k2*(4*(P0-mu1)+6*(P0*mu1-mu2)+3*(P0*mu2-mu3))
        diff = mom_eq - exact
        self.assertEqual(diff.expand(), 0)  


    def testSIS(self):
        Variables().initialise()
        
        gdd = GetStoichiometryDict(sbml_file = 'xml_files/sis.xml')
        reaction_list = gdd.getReactionList()
        
        alpha = Variables().getParameter('alpha')
        mu = Variables().getParameter('mu')
        beta = Variables().getParameter('beta')
        delta = Variables().getParameter('delta')
        
        mu10 = Variables().getMoment((1,0))
        mu01 = Variables().getMoment((0,1))
        
        mu20 = Variables().getMoment((2,0)) 
        mu11 = Variables().getMoment((1,1)) 
        mu02 = Variables().getMoment((0,2))
        
        mu30 = Variables().getMoment((3,0)) 
        mu21 = Variables().getMoment((2,1)) 
        mu12 = Variables().getMoment((1,2)) 
        mu03 = Variables().getMoment((0,3))        
        
        mom_eq = getMomentEquation((0,0), reaction_list) 
        self.assertEqual(mom_eq, 0)  
               
        #Mean 10
        mom_eq = getMomentEquation((1,0), reaction_list)
        exact = alpha - beta*mu11
        diff = mom_eq - exact
        self.assertEqual(diff, 0)  

        #Mean 01
        mom_eq = getMomentEquation((0,1), reaction_list)
        exact = delta + beta*mu11 - mu*mu01
        diff = mom_eq - exact
        self.assertEqual(diff, 0)  

        #mu20
        mom_eq = getMomentEquation((2,0), reaction_list)
        exact = alpha*(1+2*mu10) + beta*mu11-2*beta*mu21
        diff = mom_eq - exact
        self.assertEqual(diff.expand(), 0)  

        #mu11
        mom_eq = getMomentEquation((1,1), reaction_list)
        exact = beta*(mu21-mu11-mu12)-mu*mu11+delta*mu10+alpha*mu01
        diff = mom_eq - exact
        self.assertEqual(diff.expand(), 0)  
        
        #mu02
        mom_eq = getMomentEquation((0,2), reaction_list)
        exact = beta*mu11+2*beta*mu12+mu*mu01-2*mu*mu02+delta*(1+2*mu01)
        diff = mom_eq - exact
        self.assertEqual(diff.expand(), 0)  
        
    def testDummyModel(self):
        
        Variables().initialise()
        globalVariables(closure_order = 4)
        gdd = GetStoichiometryDict(sbml_file = 'xml_files/testReaction.xml')
        reaction_list = gdd.getReactionList()
       
        mom_eq = getMomentEquation((0,0,0,0,0), reaction_list) 
        self.assertEqual(mom_eq, 0)  
               
        l =Variables().getParameter('lambda')
        m = Variables().getParameter('mu')

        mu10000 = Variables().getMoment((1,0,0,0,0))
        mu01000 = Variables().getMoment((0,1,0,0,0))        
        mu00100 = Variables().getMoment((0,0,1,0,0))
        mu00010 = Variables().getMoment((0,0,0,1,0))
        mu00001 = Variables().getMoment((0,0,0,0,1))
        

        mu11000 = Variables().getMoment((1,1,0,0,0))
#        mu01000 = Variables().getMoment((0,1,0,0,0))        
#        mu00100 = Variables().getMoment((0,0,1,0,0))
#        mu00010 = Variables().getMoment((0,0,0,1,0))
#        mu00001 = Variables().getMoment((0,0,0,0,1))

                                
        #Mean 10
        print getMomentEquation((1,0,0,0,0), reaction_list)
        print getMomentEquation((0,1,0,0,0), reaction_list)
        print getMomentEquation((0,0,1,0,0), reaction_list)        
        print getMomentEquation((0,0,0,1,0), reaction_list)        
        print getMomentEquation((0,0,0,0,1), reaction_list)                        

        print getMomentEquation((2,0,0,0,0), reaction_list)
        print getMomentEquation((1,1,0,0,0), reaction_list)
        print getMomentEquation((0,2,0,0,0), reaction_list)        
        print getMomentEquation((0,0,0,0,2), reaction_list)    
#        exact = alpha - beta*mu11
#        diff = mom_eq - exact
#        self.assertEqual(diff, 0)  

        
        
        
                        
if __name__ == '__main__':
    unittest.main()




























